import os
from time import time
import argparse

MC_FILE_ROOT = os.path.join('..', '..', 'CommonFiles', 'SemanticMidFiles', 'MonteCarloEvaluation')
MINPROB = 1e-69

def monteCarloGeneratePwd(A):
    C = list()
    n = len(A)
    C.append(1.0/(A[0]*n))
    for i in range(1,n):
        if A[i] == 0.0:
            C.append(C[i - 1] + 1.0 / (MINPROB * n))
        else:
            C.append(C[i-1] + 1.0/(A[i]*n))
    return C

def listgen_dictget(ruleset,gendict):
    for home, dirs, files in os.walk(ruleset):
        for filename in files:
            category = filename.split('.')[0]
            f = open(home+'/'+filename,'r')
            tmp = {'has_P_U': False, 'key_SET': set()}
            line = f.readline()
            while line:
                line = line.strip('\n')
                if line != '':
                    linestr,lineprob = line.split('\t')
                    if len(linestr) > 4 and linestr[:4] == 'P_U_':
                        tmp['has_P_U'] = True
                        linestr = linestr.split('-')[0]
                    tmp[linestr] = float(lineprob)
                    tmp['key_SET'].add(linestr)
                line = f.readline()

            #gendict[category] = sorted(tmp.items(),key=lambda tmp:tmp[1],reverse=True)
            gendict[category] = tmp
            f.close()

def getProbFromFactor(factor, genDict):
    factorItem = factor[0]
    factorClass = factor[1]
    #print(factorItem,factorClass)
    if factorClass not in genDict:
        return 0 #need to deal with 0 problem,to find a min value

    if factorItem in genDict[factorClass]['key_SET']:
        return float(genDict[factorClass][factorItem])
    else:
        if genDict[factorClass]['has_P_U']:
            if factorItem.isdigit():
                return float(genDict[factorClass]['P_U_NUM'])
            elif factorItem.islower():
                return float(genDict[factorClass]['P_U_LOW'])
            elif factorItem.isupper():
                return float(genDict[factorClass]['P_U_UPP'])
            elif factorItem == factorItem.capitalize() and factorItem.isdigit() == False:
                return float(genDict[factorClass]['P_U_CAP'])
            else:
                return 0
        else:
            return 0

def binarySearchforA(A,pwRes):
    i = 0
    j = len(A)-1
    while i < j:
        mid = (j + i) / 2
        mid = int(mid)
        if pwRes <= A[mid]:
            i = mid + 1
        else:
            j = mid
    return mid

def splitword(testfile,seperatefile):
    wordlist = list()
    TF = open(testfile,'r')
    with open(seperatefile,'r') as f:
        line = f.readline()
        while line:
            pattern = TF.readline()
            freq1 = int(line.strip('\n').split('\t')[-1])
            freq2 = int(line.strip('\n').split('\t')[-1])
            if freq1 != freq2:
                print("These files are not justified")
            line = line.strip('\n').split('\t')[:-1]
            pattern = pattern.strip('\n').split('\t')[:-1]

            word = list()
            for i in range(len(line)):
                word.append((line[i],pattern[i]))
            word.append(freq1)
            wordlist.append(word)
            line = f.readline()
    TF.close()
    return wordlist

def generatePwd(patternlist,outputfile,wordlist,C,A,genDict):
    OF = open(outputfile,'w')
    for pwProb in wordlist:
        pwstr = ""
        pwdfreq = pwProb[-1]
        tmpPwProb = pwProb[:-1]
        for factor in tmpPwProb:
            pwstr += factor[1]
        if pwstr in patternlist:
            pwRes = patternlist[pwstr]
            for factor in tmpPwProb:
                pwRes *= getProbFromFactor(factor, genDict)
        else:
            pwRes = 0
        idx = binarySearchforA(A,pwRes)
        OF.write(str(pwdfreq) + '\t' + str(C[idx])+'\n')
    OF.close()

'''
    该文件针对SE#PCFG的蒙特卡洛随机采样结果，进行guess number的预估计算
    包含了六个参数，分别为：
    1. training_base：                   输入参数，训练库的名称
    2. target_base：                     输入参数，目标库的名称
    3. input_file_root：                 输入参数，分词结果存放的根目录
    4. training_mc_model_root：          输入参数，蒙特卡洛采样结果存放目录
    5. ruleset_name：                    输入参数，训练库使用哪个版本的分词结果。注意这里要使用Rules文件夹的全称，因为可能会有批量的实验
    6. output_root：                     输入参数，蒙特卡洛猜测模拟结果存放目录
'''
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--training_base', type=str, nargs='+', default=['CSDN_100w_20210806', 'GMAIL_100w_20210806', 'EYEEM_100w_20210806', 'FR0_100w_20210806'])
    parser.add_argument('--target_base', type=str, nargs='+', default=['17173_100w_20210806', 'Twitter_100w_20210806', '8Fit_100w_20210806', 'FR1_100w_20210806'])
    parser.add_argument('--input_file_root', default=os.path.join('..', 'Results', 'Result_20221029'))
    parser.add_argument('--training_mc_model_root', default=os.path.join(MC_FILE_ROOT, 'MonteCarlo_20221213', 'mc_models'))
    parser.add_argument('--ruleset_name', default='Rules')
    parser.add_argument('--output_root', default=os.path.join(MC_FILE_ROOT, 'MonteCarlo_20221213', 'mc_evaluations'))

    args = parser.parse_args()
    training_base_list = args.training_base
    target_base_list = args.target_base
    input_file_root = args.input_file_root
    training_mc_model_root = args.training_mc_model_root
    output_root = args.output_root
    ruleset_name = args.ruleset_name

    if not os.path.exists(output_root):
        os.makedirs(output_root)

    for this_train in training_base_list:
        for this_target in target_base_list:
            begin_time = time()
            patternfile = os.path.join(input_file_root, this_train, '%s%s' % (this_train, '_genstruct'))
            seperatefile = os.path.join(input_file_root, this_target, '%s%s' % (this_target, '_seperate'))
            samplefile = os.path.join(training_mc_model_root, '%s%s' % (this_train, '_MC_model'))
            ruleset = os.path.join(input_file_root, this_train, ruleset_name)
            testfile = os.path.join(input_file_root, this_target, '%s%s' % (this_target, '_struct'))
            outputfile = os.path.join(output_root, 'raw_mc_evaluation_%s_to_%s' % (this_train, this_target))

            patternlist = dict()
            with open(patternfile, 'r') as f:
                pw = f.readline()
                while pw:
                    pw = pw.strip('\n')
                    pwstr = ''.join(pw.split('\t')[:-1])
                    pwprob = float(pw.split('\t')[-1])
                    patternlist[pwstr] = pwprob
                    pw = f.readline()
            genDict = dict()
            listgen_dictget(ruleset, genDict)
            A = list()
            with open(samplefile, 'r') as f:
                pw = f.readline()
                while pw:
                    pw = pw.strip('\n')
                    A.append(float(pw))
                    pw = f.readline()
            C = monteCarloGeneratePwd(A)
            wordlist = splitword(testfile, seperatefile)
            print(len(wordlist))
            generatePwd(patternlist, outputfile, wordlist, C, A, genDict)
            print('totally cost time:', time()-begin_time)

if __name__ == '__main__':
    main()
